import cv2
import os
import os.path as osp
import datetime

run_time = '_'.join(str(datetime.datetime.now()).split(':'))

BASEDIR = osp.dirname(osp.abspath(__file__))
outdir = osp.join(BASEDIR, 'output' + run_time)
if not osp.exists(outdir):
    os.makedirs(outdir)

detector = cv2.CascadeClassifier("./haarcascade_frontalface_alt.xml")

imgnames = [
    fname for fname in os.listdir(BASEDIR)
    if fname.split('.')[-1] in ['png', 'jpg', 'jpeg', 'bmp', 'webp', 'tif']
]

for ii, imgname in enumerate(imgnames):
    path_img = osp.join(BASEDIR, imgname)
    # print(path_img, 'processing...')

    image = cv2.imread(path_img)
    rects = detector.detectMultiScale(
        image,
        scaleFactor=1.1,  # 严格大于1，越小，出的框越多
        minNeighbors=2,  # 越小，出的框越多
        minSize=(30, 30))

    for i, (x, y, w, h) in enumerate(rects):
        try:
            img_patch = image[y - h // 4:y + h + h // 4,
                              x - w // 4:x + w + w // 4, :]
            out_path = osp.join(outdir,
                                f"{imgname.split('.')[-2]}_face_{ii}_{i}.jpg")
            cv2.imwrite(out_path, img_patch)
        except:
            img_patch = image[y:y + h, x:x + w, :]
            out_path = osp.join(outdir,
                                f"{imgname.split('.')[-2]}_face_{ii}_{i}.jpg")
            cv2.imwrite(out_path, img_patch)
        print(imgname, 'done.')
